{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {func}`tvm.relay.qnn.op.softmax`\n",
    "\n",
    "源码：``tvm/src/relay/qnn/op/softmax.cc``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm import relay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mSignature:\u001b[0m\n",
      "\u001b[0mrelay\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mqnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mscale\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mzero_point\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0moutput_scale\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0moutput_zero_point\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m <no docstring>\n",
      "\u001b[0;31mSource:\u001b[0m   \n",
      "\u001b[0;32mdef\u001b[0m \u001b[0msoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mzero_point\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_scale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_zero_point\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mreturn\u001b[0m \u001b[0m_make\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msoftmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mzero_point\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_scale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_zero_point\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFile:\u001b[0m      /media/pc/data/lxw/ai/tvm/python/tvm/relay/qnn/op/qnn.py\n",
      "\u001b[0;31mType:\u001b[0m      function"
     ]
    }
   ],
   "source": [
    "relay.qnn.op.softmax??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Not equal to tolerance rtol=1e-07, atol=1\n",
      "\n",
      "Mismatched elements: 5 / 50 (10%)\n",
      "Max absolute difference: 5\n",
      "Max relative difference: 0.33333333\n",
      " x: array([[-128, -128, -128, -128, -128, -128, -128, -128, -128,  126],\n",
      "       [-128, -128, -128, -128, -128, -128, -128, -124,  -96,   88],\n",
      "       [-128, -128, -128, -128, -128, -128, -128, -128, -120,  118],...\n",
      " y: array([[-128, -128, -128, -128, -128, -128, -128, -128, -128,  127],\n",
      "       [-128, -128, -128, -128, -128, -128, -128, -123,  -98,   93],\n",
      "       [-128, -128, -128, -128, -128, -128, -128, -128, -120,  120],...\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relay\n",
    "\n",
    "is_sorted = lambda a: np.all(a[:-1] <= a[1:])\n",
    "\n",
    "shape = [5, 10]\n",
    "scale = 0.2\n",
    "x_ = relay.var(\"x\", shape=shape, dtype=\"int8\")\n",
    "x = relay.qnn.op.dequantize(x_, relay.const(scale), relay.const(0))\n",
    "op = relay.op.nn.softmax(x, axis=1)\n",
    "op = relay.qnn.op.quantize(\n",
    "    op, relay.const(1.0 / 256.0), relay.const(-128), out_dtype=\"int8\"\n",
    ")\n",
    "\n",
    "x_np = np.random.randint(-128, 127, size=shape, dtype=\"int8\")\n",
    "x_np = np.sort(x_np)\n",
    "args = [x_np]\n",
    "\n",
    "mod = tvm.IRModule.from_expr(op)\n",
    "mod = tvm.relay.transform.InferType()(mod)\n",
    "mod_int = tvm.relay.transform.FakeQuantizationToInteger(\n",
    "    hard_fail=True, optional_qnn_ops=[\"nn.softmax\"]\n",
    ")(mod)\n",
    "assert not tvm.ir.structural_equal(mod, mod_int)\n",
    "result = (\n",
    "    relay.create_executor(\"vm\", mod=mod, device=tvm.cpu(), target=\"llvm\")\n",
    "    .evaluate()(*args)\n",
    "    .numpy()\n",
    ")\n",
    "result_int = (\n",
    "    relay.create_executor(\"vm\", mod=mod_int, device=tvm.cpu(), target=\"llvm\")\n",
    "    .evaluate()(*args)\n",
    "    .numpy()\n",
    ")\n",
    "\n",
    "# Check at least the softmax output is in ascending order,\n",
    "# since it is difficult to use allclose due to not-so-good accuracy.\n",
    "for qdq, qop in zip(result, result_int):\n",
    "    assert is_sorted(qdq)\n",
    "    assert is_sorted(qop)\n",
    "\n",
    "try:\n",
    "    np.testing.assert_allclose(result_int, result, atol=1)\n",
    "except AssertionError as e:\n",
    "    # To see the difference\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] {\n",
       "  qnn<span style=\"color: #AA22FF; font-weight: bold\">.</span>softmax(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #008000\">0.2</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #008000\">0.00390625</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">-</span><span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, axis<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod_int.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> max(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, axis<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], keepdims<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">1</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> right_shift(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> right_shift(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #008000\">5</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">2.14748e+09</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2.14748e+09</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> negative(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> divide(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">20</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> negative(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">14</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">15</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">14</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> right_shift(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">15</span>, <span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">17</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> max(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, axis<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], keepdims<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">1</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">18</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">16</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">19</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">17</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">20</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> left_shift(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">18</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">19</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">21</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> sum(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">20</span>, axis<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], keepdims<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">1</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">22</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> divide(<span style=\"color: #008000\">1073741824</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">21</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">1</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">23</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">22</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">20</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">24</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> right_shift(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">23</span>, <span style=\"color: #008000\">23</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">25</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">24</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">26</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">-</span><span style=\"color: #008000\">128</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">27</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fixed_point_multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">25</span>, multiplier<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1073741824</span>, shift<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">28</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">26</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">27</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">29</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">28</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">128</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">127</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">29</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with tvm.target.Target(\"llvm\"):\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        run_mod = relay.qnn.transform.Legalize()(mod_int)\n",
    "        run_mod = relay.qnn.transform.CanonicalizeOps()(run_mod)\n",
    "run_mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "/*\n",
    " * \\brief Canonicalizes the QNN softmax op.\n",
    " * \\param attrs The Softmax attrs.\n",
    " * \\param new_args The new mutated args to the call node.\n",
    " * \\param arg_types The types of input and output.\n",
    " * \\return The sequence of Relay ops for softmax op.\n",
    " * \\note This op is highly experimental and sometimes lacks accuracy.\n",
    " *       Be aware that the input scale must be in the range of 0 to 1.\n",
    " */\n",
    "Expr QnnSoftmaxCanonicalize(const Attrs& attrs, const Array<Expr>& new_args,\n",
    "                            const Array<tvm::relay::Type>& arg_types) {\n",
    "  // Expected: input, scale, zero_point, output_scale, output_zero_point\n",
    "  ICHECK_EQ(new_args.size(), 5);\n",
    "\n",
    "  const auto const_i32 = [&](int32_t val) { return MakeConstantScalar(DataType::Int(32), val); };\n",
    "  const auto const_f32 = [&](float val) { return MakeConstantScalar(DataType::Float(32), val); };\n",
    "\n",
    "  const auto const_input_scale = new_args[1].as<ConstantNode>();\n",
    "  ICHECK(const_input_scale) << \"Input scale should be constant.\";\n",
    "  ICHECK(const_input_scale->is_scalar()) << \"Input scale should be scalar.\";\n",
    "  const float input_scale = static_cast<float*>(const_input_scale->data->data)[0];\n",
    "  ICHECK(input_scale <= 1.f) << \"Input scale should be less than or equal to 1.\";\n",
    "\n",
    "  const Expr input_zero_point = new_args[2];\n",
    "  const Expr output_scale = new_args[3];\n",
    "  const Expr output_zero_point = new_args[4];\n",
    "  const int axis = attrs.as<SoftmaxAttrs>()->axis;\n",
    "\n",
    "  // Refer to the Algorithm 1 in https://arxiv.org/pdf/2207.01405.pdf\n",
    "\n",
    "  const Expr quantized_data = Subtract(Cast(new_args[0], DataType::Int(32)), input_zero_point);\n",
    "\n",
    "  const Expr x_0 = ConvertDtype(const_f32(std::round(1.f / input_scale)), DataType::Int(32));\n",
    "  const Expr max = Max(quantized_data, {axis}, true, false);\n",
    "  const Expr x = Subtract(quantized_data, max);\n",
    "\n",
    "  const int m = 30;\n",
    "  const int bits = 8;\n",
    "  const Expr x_p = Subtract(Add(x, RightShift(x, const_i32(1))), RightShift(x, const_i32(4)));\n",
    "  const Expr q = Clip(Divide(x_p, Negative(x_0)), 0, 20);\n",
    "  const Expr max_q = Max(q, {axis}, true, false);\n",
    "  const Expr r = Subtract(x_p, Multiply(q, Negative(x_0)));\n",
    "  const Expr x_b = Add(RightShift(r, const_i32(1)), x_0);\n",
    "  const Expr exps = LeftShift(x_b, Subtract(max_q, q));\n",
    "  const Expr sums = Sum(exps, {axis}, true, false);\n",
    "  const Expr output =\n",
    "      RightShift(Multiply(Divide(const_i32(1 << m), sums), exps), const_i32(m - (bits - 1)));\n",
    "  const Expr requantized = Requantize(output, arg_types[0].as<TensorTypeNode>()->shape,\n",
    "                                      const_f32(1.f / (1 << (bits - 1))), const_i32(0),\n",
    "                                      output_scale, output_zero_point, DataType::Int(bits), 0);\n",
    "\n",
    "  return requantized;\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这段代码是一个用于规范化 QNN softmax 操作的函数。它接受三个参数：`attrs`（Softmax 属性）、`new_args`（新的调用节点参数）和 `arg_types`（输入和输出的类型）。该函数返回  Relay 算子序列，用于执行 softmax 运算。\n",
    "\n",
    "该函数首先检查输入参数的数量是否正确，然后从 `new_args` 中提取出输入、缩放因子、零点、输出缩放因子和输出零点等参数。接着，它使用这些参数计算出量化数据，并按照[算法 1](https://arxiv.org/pdf/2207.01405.pdf) 进行计算。最后，它将结果重新量化并返回。\n",
    "\n",
    "需要注意的是，这个函数是高度实验性的，有时可能缺乏准确性。此外，输入缩放因子必须在 0 到 1 之间。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312x",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.1.-1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
